import csv


def clean_student_faculty_ratio(ratio_str):
    """Extract the numeric value from student-faculty ratio string (e.g., '1:15' -> 15)"""
    if isinstance(ratio_str, str) and ":" in ratio_str:
        try:
            return int(ratio_str.split(":")[1])
        except (ValueError, IndexError):
            return None
    return None


def calculate_student_count(ratio_value, staff_count):
    """Calculate estimated student count using ratio and staff count"""
    if ratio_value and staff_count:
        return ratio_value * staff_count
    return None


# Load and process data
data = []
with open("baoshan-schools-2024.csv", mode="r", encoding="utf-8") as file:
    reader = csv.reader(file)
    for row in reader:
        if row and row[0] != "" and row[2] == "公办":
            data.append(row)

# Process data
processed_data = []
header = [
    "serial_no",
    "school_name",
    "total_staff_count",
    "full_time_teacher_count",
    "estimated_student_count"
]

processed_data.append(header)  # Add header row

for row in data[1:]:  # Skip original header
    try:
        serial_no = int(row[0])
        school_name = row[1]
        total_staff = int(row[14])
        full_time_teachers = int(row[15])
        ratio_str = row[16]

        ratio_value = clean_student_faculty_ratio(ratio_str)
        if ratio_value is None:
            continue  # Skip rows without valid ratio

        estimated_students = calculate_student_count(ratio_value, total_staff)
        if estimated_students is None:
            continue

        processed_data.append([
            serial_no,
            school_name,
            total_staff,
            full_time_teachers,
            estimated_students
        ])
    except (ValueError, IndexError):
        continue  # Skip rows with invalid data

# Save output
with open("clean-baoshan-schools-2024.csv", mode="w", encoding="utf-8", newline="") as file:
    writer = csv.writer(file)
    writer.writerows(processed_data)
